from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts.prompt import PromptTemplate
from langchain.chains import LLMChain

import os

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
# 定义模型
llm = Tongyi()

# 定义两个chain
synopsis_template = """你是一名写作家. 给定游戏的标题,以及类型，然后你的工作就是为该标题结合类型写一个概要.
标题: {title}
类型: {type}
写作家: 这是以上游戏的简介:"""
synopsis_prompt_template = PromptTemplate(input_variables=["title", "type"], template=synopsis_template)
synopsis_chain = LLMChain(llm=llm, prompt=synopsis_prompt_template, output_key='content')

commentator_template = """你是一名评论员. 给定游戏的简介,你的工作就是为该游戏，给出评分（1-10）.
游戏简介概要: {content}
评论员: 这是对于该游戏的评论和评分:"""
commentator_prompt_template = PromptTemplate(input_variables=["content"], template=commentator_template)
commentator_chain = LLMChain(llm=llm, prompt=commentator_prompt_template, output_key='comment')

# 组装
from langchain.chains import SequentialChain
overall_chain = SequentialChain(chains=[synopsis_chain, commentator_chain],
                                input_variables=['title', 'type'],
                                output_variables=["content", "comment"],
                                verbose=True)

res = overall_chain.invoke({'title': '奥特曼大战怪兽', 'type': '恐怖'})
print(res)
